#!/usr/bin/python

import numpy

#26 - ' ' 32
#27 - .   46
#28 - ,   44
#29 - !   33
#30 - ?   63

#this class implements a Hill Cipher, a type of polygraphic cipher
class PolygramCipher:
    def __init__(self,n):
        self.n = int(n)
        #create the matrix keys
        if self.n == 2:
            A = [[5,8],[17,3]]
            B = [[9,2],[1,15]]
        elif self.n == 3:
            A = [[6,24,1],[13,16,10],[20,17,15]]
            B = [[8,5,10],[21,8,21],[21,12,8]]
        elif self.n == 4:
            A = [[25,29,20,23],[30,16,9,10],[15,2,5,3],[4,30,10,19]]
            B = [[2,20,26,27],[4,19,22,29],[1,18,16,8],[27,2,15,8]]
        else:
            A = []
            B = []
        self.key = numpy.array(A)
        self.invkey = numpy.array(B)

    #convert a string into a numerical matrix
    def convertText(self,message):
        m = []
        for c in message:
            if c == ' ':
                x = 26
            elif c == '.':
                x = 27
            elif c == ',':
                x = 28
            elif c == '!':
                x = 29
            elif c == '?':
                x = 30
            else:
                x = ord(c)-97
            m.extend([[x]])            #add an element to the array
        return numpy.array(m)

    #convert a numerical matric back into a string
    def convertMatrix(self,matrix):
        m = ''
        for e in matrix:
            for c in e:
                if c == 26:
                    m += chr(32)
                elif c == 27:
                    m += chr(46)
                elif c == 28:
                    m += chr(44)
                elif c == 29:
                    m += chr(33)
                elif c == 30:
                    m += chr(63)
                else:
                    m += chr(int(c+97))
        return m

    #encrypts a message using a hill cipher
    def encrypt(self,message):
        ciphertext = ''
        j = 0
        while j < len(message):
            m = ''
            if j+self.n >= len(message):
                if j == len(message)-1:
                    m = message[j]
                else:
                    m = message[j:len(message)]
                    #print m
                i = self.n+j-len(message)
                while i > 0:
                    i -= 1
                    m += ' '
            else:
                m = message[j:j+self.n]
            ciphertext += self.e(m)
            j = j+self.n
        return ciphertext

    def e(self,message):
        #convert message to a matrix
        x = self.convertText(message)
        #matrix multiple the key matrix by the message matrix
        ciphertext = numpy.dot(self.key,x).tolist()
        #mod answer by 26
        for i in range(0,len(ciphertext)):
            y = ciphertext[i]
            for j in range(0,len(y)):
                y[j] = y[j] % 31
                #y[j] = y[j] % 26
            ciphertext[i] = y
        #convert back to text and return ciphertext
        return self.convertMatrix(ciphertext)

    def decrypt(self,ciphertext):
        message = ''
        j = 0
        while j < len(ciphertext):
            m = ciphertext[j:j+self.n]
            message += self.d(m)
            j += self.n
        return message

    #decrypts a message using a hill cipher
    def d(self,ciphertext):
        #convert ciphertext to matrix
        x = self.convertText(ciphertext)
        #matrix multple the inverse key matrix by the message matrix
        message = numpy.dot(self.invkey,x).tolist()
       #mod answer by 26
        for i in range(0,len(message)):
            y = message[i]
            for j in range(0,len(y)):
                y[j] = y[j] % 31
                #y[j] = y[j] % 26
            message[i] = y
        #convert back to text and return unencrypted message
        return self.convertMatrix(message)
